from codes.util.data_util import get_shape_list
from codes.util.model_util import conv2d, conv2d_transform, batch_norm, relu, lrelu, element_wise_sum, pixel_shuffer, full_connected

def residual_block(input_data, output_shape):
    conv_data=conv2d(input_data, output_shape, 3, 3, 1, 1)
    norm_data=batch_norm('batch_norm',conv_data)
    relu_data=relu(norm_data)
    conv_data=conv2d(relu_data, output_shape, 3, 3, 1, 1)
    norm_data = batch_norm('batch_norm', conv_data)
    return element_wise_sum(input_data, norm_data)

def discriminate_block(input_data,output_shape_item,d_h,d_w,is_norm=False):
    output_shape_list = get_shape_list(input_data).insert(output_shape_item)
    conv_data = conv2d(input_data, output_shape_list, 3, 3, d_h, d_w)
    lrelu_data = lrelu(conv_data)

    if is_norm:
        return batch_norm('batch_norm', lrelu_data)

    return lrelu_data